import os
import torch
from torch import nn
import my_train
import plot_2D
import plot_3D
import plot_eg2
import plot_eg6
from NN import ann, cdinn2
import time
from prob import *
from generate_data import *
import global_verify
from verify_dccp import verify_ccp
from verify_mip_relu import verify_barrier_mip_relu


torch.set_default_dtype(torch.float64)
torch.set_default_tensor_type(torch.DoubleTensor)

model_dir = "./model/"
log_dir = "./log/"

if not os.path.exists(model_dir):
    os.mkdir(model_dir)
if not os.path.exists(log_dir):
    os.mkdir(log_dir)

if USE_CDINN:
    model = cdinn2.gen_nn()
else:
    model = ann.gen_nn()

# model.load_state_dict(torch.load(model_dir + "cdinn_c1_1_5_0_0.05_1_pre.pt"), strict=True)

model_name = ("cdinn_" if USE_CDINN else "ann_") + case_name + f"_{N_H}_{D_H}_{alpha}_{time_step}_{piece_len}_train.pt"
# model.load_state_dict(torch.load(model_dir + model_name), strict=True)
# model.load_state_dict(torch.load('./model/cdinn_eg4_retry_1_20_0.05_0.1_1_train.pt'), strict=True)
# model.load_state_dict(torch.load('./model/cdinn_zhrC13_1_20_0.1_0.1_1_train.pt'), strict=True)
# model.load_state_dict(torch.load('./model/cdinn_zhrC14_1_20_0.1_0.1_1_train.pt'), strict=True)
# model.load_state_dict(torch.load('./model/cdinn_c1_1_5_0.1_0.1_1_train.pt'), strict=True)
# model.load_state_dict(torch.load('./model/cdinn_zhrR9_woRandomGen_1_20_0.1_0.1_0.3_train.pt'), strict=True)
# model.load_state_dict(torch.load('./model/cdinn_zhrc1_1_20_0.1_0.1_1_train.pt'), strict=True)
# model.load_state_dict(torch.load('./model/cdinn_c9_1_5_0.1_0.1_1_train.pt'), strict=True)

train_veri_succ = False
train_succ = False

print("generating data...")
time_start_data = time.time()
if case_name == 'eg2':
    batches_init, batches_unsafe, batches_domain = generate_data_eg2()
elif case_name == 'eg6':
    batches_init, batches_unsafe, batches_domain = generate_data_eg6()
else:
    batches_init, batches_unsafe, batches_domain = generate_data()
time_end_data = time.time()
print("Data generate time cost:  " + str(time_end_data - time_start_data))
# print("first_addcex:", len(batches_domain))

if 1==1:

    for restart_n in range(RESTART_NUMS):
        print("\n=====================================")
        print(f"{restart_n + 1} times restart the eg!\n")
        with torch.no_grad():
            for _, p in model.named_parameters():
                nn.init.normal_(p)
        # torch.save(model.state_dict(), model_dir + model_name_pre)

        train_veri_succ = False
        train_succ = False

        rnd = 0
        time_train_sum = 0
        time_verify_sum = 0

        log_name_back = ("ccp_" if USE_CDINN else "fc_") + case_name + f"_{N_H}_{D_H}_{alpha}_{time_step}_{piece_len}_time.txt"
        log_name = ("ccp_" if USE_CDINN else "fc_") + case_name + "_time.txt"
        fw_back = open(log_dir + log_name_back, "w+", encoding="utf-8")
        fw = open(log_dir + log_name, "w+", encoding="utf-8")
        fw.write(case_name + f"_{N_H}_{D_H}_{alpha}_{time_step}_{piece_len}_time.txt\n\n")

        fw.write("restart: " + str(restart_n) + "\n")
        fw_back.write("restart: " + str(restart_n) + "\n")
        fw.write("piece_num: " + str(piece_num) + "\n")
        fw_back.write("piece_num: " + str(piece_num) + "\n")
        fw.write("init_min: " + str(init_min) + "\n")
        fw_back.write("init_min: " + str(init_min) + "\n")
        fw.write("init_max: " + str(init_max) + "\n")
        fw_back.write("init_max: " + str(init_max) + "\n")
        fw.write("unsafe_min: " + str(unsafe_min) + "\n")
        fw_back.write("unsafe_min: " + str(unsafe_min) + "\n")
        fw.write("unsafe_max: " + str(unsafe_max) + "\n")
        fw_back.write("unsafe_max: " + str(unsafe_max) + "\n")

        while True:
            rnd = rnd + 1
            fw_back.write(f"\n第{rnd}轮：\n")
            fw.write(f"\n第{rnd}轮：\n")
            time_start_train = time.time()
            model, train_succ, verify_error = my_train.train(model, batches_init, batches_unsafe, batches_domain)
            time_end_train = time.time()
            time_train = time_end_train - time_start_train
            time_train_sum += time_train

            fw_back.write(f"本轮训练时间：{time_train}\n")
            fw.write(f"本轮训练时间：{time_train}\n")
            if (not train_succ) or verify_error:
            # if not train_succ:
                # train_veri_succ = True
                break
            else:
                # torch.save(model.state_dict(), model_dir + 'cex_train.pt')
                time_start_verify = time.time()
                if USE_CDINN:
                    # cex_init_curr, cex_unsafe_curr, cex_domain_curr = verify_barrier_dccp.verif(model, piece_len)
                    # cex_init_curr, cex_unsafe_curr, cex_domain_curr = verify_ccp.verif(model, piece_len)
                    cex_init_curr, cex_unsafe_curr, cex_domain_curr = global_verify.verif(model, piece_len)
                else:
                    cex_init_curr, cex_unsafe_curr, cex_domain_curr = verify_barrier_mip_relu.verify_relu(model, piece_len)
                time_end_verify = time.time()

                time_verify = time_end_verify - time_start_verify
                time_verify_sum += time_verify
                fw_back.write(f"本轮验证时间：{time_verify}\n\n")
                fw.write(f"本轮验证时间：{time_verify}\n\n")

                if len(cex_init_curr) == 0 and len(cex_unsafe_curr) == 0 and len(cex_domain_curr) == 0:
                    train_veri_succ = True
                    break

                # if set_time_limit and time_train_sum > time_limit:
                #     break
                batches_init, batches_unsafe, batches_domain = add_cexes(batches_init, cex_init_curr, batches_unsafe, cex_unsafe_curr, batches_domain, cex_domain_curr)
                # print("after_addcex:", len(batches_domain))
        # if verify_error:
        #     break

        if train_succ and train_veri_succ:
            break
        # if set_time_limit and time_train_sum > time_limit:
        #     continue

    fw_back.write(f"\n\n本次训练总时间: {time_train_sum}\n")
    fw.write(f"\n\n本次训练总时间: {time_train_sum}\n")
    fw_back.write(f"本次验证总时间: {time_verify_sum}\n\n")
    fw.write(f"本次验证总时间: {time_verify_sum}\n\n")
    print(f"本次训练总时间: {time_train_sum}\n")
    print(f"本次验证总时间: {time_verify_sum}\n")
    print(f"rnd:", rnd)

    fw_back.write(f"train_succ: {train_succ}\n")
    fw_back.write(f"train_veri_succ: {train_veri_succ}\n")
    fw.write(f"train_succ: {train_succ}\n")
    fw.write(f"train_veri_succ: {train_veri_succ}\n")
    fw_back.close()
    fw.close()


# print("\nData generation totally costs:", time_end_data - time_start_data)
# print("Training totally costs:", time_end_train - time_start_train)

# save model
if train_succ and train_veri_succ:
    torch.save(model.state_dict(), model_dir + model_name)

print("train_succ: ", train_succ)
print("train_veri_succ: ", train_veri_succ)

# plot
# if case_name == 'eg2':
#     plot_eg2.simulation(model)
# elif case_name == 'eg6':
#     plot_eg6.simulation(model)
# elif DIM == 2:
#     plot_2D.simulation(model)
# elif DIM == 3:
#     plot_3D.simulation(model)


